from calc_manager.plugins_inner.bdf_reader_inner import BDFReader
from calc_manager.plugins_inner.op2_reader_inner import OP2Reader
from calc_manager.plugins_inner.mat_csv_inner import MatCSVReader
from calc_manager.plugins_inner.bdf_transformer_inner import BDFTransformer
from pydantic import BaseModel
from typing import List, Dict

class Parameter(BaseModel):
    rough_eids: List[float]
    refine_eids: List[float]

class Result(BaseModel):
    margin_output: List[Dict]
    danger_infos: Dict[str, List[float]]
    rough_eids: List[float]


def second(param: Parameter, bdf_reader: BDFReader,  mat_csv_reader: MatCSVReader, bdf_transformer: BDFTransformer, op2_reader: OP2Reader) :
    print("step2-----------------> start")
    rough_eids = param.rough_eids
    refine_eids = param.refine_eids
    rough_eids_set = set(rough_eids) # Convert to set for O(1) lookup

    results = {}
    subcase_id = op2_reader.get_subcase_id()
    result_components = ["stress.ctria3_stress", "stress.cquad4_stress"]
    for comp in result_components:
        print(f"step2-----------------> Getting results for component {comp} and subcase {subcase_id}")
        stress = op2_reader.get_result_from_comp_and_subcase(comp, subcase_id)
        print(f"step2-----------------> Got stress results with shape: {stress.data.shape}")
        element_node = stress.element_node
        von_mises = stress.von_mises()
        for index, item in enumerate(element_node):
            if index % 2 == 0:
                eid = item[0]
                # 只保存非rough_eids的结果 , 过滤掉自然网格
                if eid not in rough_eids_set: # 使用set进行O(1)的查找
                    results[eid] = von_mises[0][index : index + 2].tolist()

    print(f"step2-----------------> results: complete")

        #for index, item in enumerate(element_node):
        #    if index % 2 == 0:
        #        results[item[0]] = von_mises[0][index : index + 2].tolist()

    ## 删除自然网格的结果(无需关注)
    #for eid in rough_eids:
    #    results.pop(eid, None)

    """
    1. 失效单元eid
    2. 失效单元eid对应的cid,对应cname
    3. 根据cname获取到g70name,再分组
    4. 第二张图: 根据g70name分组, 补充当前失效cnames包含的所有eids
    6. 第三张图: 根据g70name分组, 补充当前失效g70下所有的cnames的所有eids
    """
    ## TODO 理论上需要遍历所有的bdf文件获取所有数据,目前暂时指定文件名
    #refine_bdf_path = r"D:\SF-CAE\model\S20.bdf"  # 需要替换为真实的精细网格bdf文件
    #bdf_transformer = BDFTransformer(refine_bdf_path)
    print(f"step2-----------------> get comps_info")
    comps_info = bdf_transformer.get_comps_info()

    # 获取所有eid对应的安全系数以及失效单元eid
    # danger_infos包含g70_name和对应的cids关系
    # 组织层级 gname-->cname-->eid
    print(f"step2-----------------> get danger_infos")
    danger_infos = {}
    for refine_eid in refine_eids:
        element = bdf_reader.get_element(refine_eid)
        mat_id = element.pid_ref.material_ids[0]
        # 根据mat_id获取材料许用应力数据
        material_ultimate_stress = mat_csv_reader.get_mat_from_csv(mat_id)
        if refine_eid not in results:
            print(f"refine_eid: {refine_eid} not in results, results keys: {results.keys()}")
            raise Exception(f"refine_eid: {refine_eid} not in results")

        element_stress = results[refine_eid]
        element_stress.append(material_ultimate_stress / 3)
        coffe = material_ultimate_stress / max(element_stress) - 1
        if coffe < 0:
            # 获取cid
            cid = bdf_transformer.get_comp_by_eid(refine_eid)
            for comp_info in comps_info:
                if cid == comp_info["cid"]:
                    cname = comp_info["cname"]
                    # 上航提供的名称规则
                    g70_name = cname.split("_")[0]
                    break
            if g70_name not in danger_infos:
                danger_infos[g70_name] = [cid]
            elif cid not in danger_infos[g70_name]:
                danger_infos[g70_name].append(cid)
        results[refine_eid].append(coffe)

    print(f"step2-----------------> get danger_infos: complete")

    """
    构造最终裕度图需要的数据
    1. output长度代表最终报表的行数
    2. gname代表表格第一列的G70编号结果
    3. 基于eids和coffe三维负责渲染数据
    4. num代表表格第三列的失效单元数
    """
    print(f"step2-----------------> get margin_output")
    margin_output = []
    for k, v in danger_infos.items():
        eids = []
        for cid in v:
            eids.extend(bdf_transformer.get_comp_eids_ref()[cid])
        coffe = [results[eid][-1] for eid in eids]
        num = sum(1 for i in coffe if i < 0)
        margin_output.append(
            {
                "gname": k,
                "eids": eids,
                "coffe": coffe,
                "num": num,
            }
        )
    print(f"step2-----------------> get margin_output: complete")
    res =  Result(margin_output=margin_output, danger_infos=danger_infos, rough_eids=rough_eids)
    print("step2-----------------> end")
    return res
